import io
from os import path, makedirs
from re import search
from time import time
from requests import Session
from PIL import Image

import pytesseract

# JAcount username and password
USERNAME = ''
PASSWORD = ''

# number of images to fetch
FETCH_COUNT = 10000
# ratio of training set
TRAIN_RATIO = 0.9
# get captcha code
CAPTURE_MODE = 1 # 0: manual, 1: pytesseract

LOGIN_URL = 'https://jaccount.sjtu.edu.cn/oauth2/authorize?response_type=code&scope=profile&client_id=mB5nKHqC00MusWAgnqSF&redirect_uri=https://sports.sjtu.edu.cn/oauth2Login'


train_data_path = path.join(path.dirname(path.abspath(__file__)), 'data/train')
test_data_path = path.join(path.dirname(path.abspath(__file__)), 'data/test')
if not path.exists(train_data_path):
    makedirs(train_data_path)
if not path.exists(test_data_path):
    makedirs(test_data_path)


def login(session, username, password):
    res = session.get(LOGIN_URL).text

    sid = search(r'sid: "(.*?)"', res).group(1)
    returl = search(r'returl:"(.*?)"', res).group(1)
    se = search(r'se: "(.*?)"', res).group(1)
    client = search(r'client: "(.*?)"', res).group(1)
    uuid = search(r'captcha\?uuid=(.*?)&t=', res).group(1)

    captcha_id = search(r'img.src = \'captcha\?(.*)\'', res).group(1) + str(int(time() * 1000))
    captcha_url = 'https://jaccount.sjtu.edu.cn/jaccount/captcha?' + captcha_id
    captcha = session.get(captcha_url, headers={'Referer': 'https://jaccount.sjtu.edu.cn'})

    captcha_img = Image.open(io.BytesIO(captcha.content))
    # get the captcha code
    if CAPTURE_MODE == 0: # manual
        with open('captcha.jpeg', 'wb') as f:
            f.write(captcha.content)
        captcha_img = Image.open('captcha.jpeg')
        captcha_code = input('请输入验证码：')
    elif CAPTURE_MODE == 1: # pytesseract
        captcha_code = pytesseract.image_to_string(captcha_img)
        captcha_code = captcha_code.strip()
        print(captcha_code)

    data = {
        'sid': sid,
        'client': client,
        'returl': returl,
        'se': se,
        'v': '',
        'uuid': uuid,
        'user': username,
        'pass': password,
        'captcha': captcha_code,
    }
        
    res = session.post('https://jaccount.sjtu.edu.cn/jaccount/ulogin', data=data)
    # with open('res.html', 'w', encoding='UTF-8') as f:
    #     f.write(res.text)
    if 'SJTU Single Sign On' in res.text:
        raise Exception('login error')
    return captcha_code, captcha_img


def main():
    # split training and testing data
    train_data_len = int(FETCH_COUNT * TRAIN_RATIO)
    test_data_len = FETCH_COUNT - train_data_len

    ## 1. collect training data
    # read exsit number of images
    tain_label_file = path.join(train_data_path, 'labels.txt')
    i = 0
    if path.exists(tain_label_file):
        print('train data exists, continue collecting...')
        label_file = open(tain_label_file, 'a')
        label_file.write('\n')
        with open(tain_label_file, 'r') as _:
            i = len(_.readlines())
    else:
        label_file = open(tain_label_file, 'w')
        
    while i < train_data_len:
        session = Session()
        # login
        try:
            captcha_code, captcha_img = login(session, USERNAME, PASSWORD)
        except Exception as e:
            print('[ERROR]: ', e)
            continue

        print(f'train[{i+1}/{train_data_len}] captcha code: {captcha_code}')
        # save captcha image
        captcha_img.save(path.join(train_data_path, f'{i}.png'))
        # save captcha code
        label_file.write(f'{captcha_code}\n')
        session.close()
        i += 1
    label_file.close()
    
    ## 2. collect testing data
    # read exsit number of images
    test_label_file = path.join(test_data_path, 'labels.txt')
    i = 0
    if path.exists(test_label_file):
        print('test data exists, continue collecting...')
        label_file = open(test_label_file, 'a')
        label_file.write('\n')
        with open(test_label_file, 'r') as _:
            i = len(_.readlines())
    else:
        label_file = open(test_label_file, 'w')

    while i < test_data_len:
        session = Session()
        # login
        try:
            captcha_code, captcha_img = login(session, USERNAME, PASSWORD)
        except Exception as e:
            print('[ERROR]: ', e)
            continue

        print(f'test[{i+1}/{test_data_len}] captcha code: {captcha_code}')
        # save captcha image
        captcha_img.save(path.join(test_data_path, f'{i}.png'))
        # save captcha code
        label_file.write(f'{captcha_code}\n')
        session.close()
        i += 1
    label_file.close()

if __name__ == '__main__':
    main()
